import sys
import cv2
import os
import numpy as np
from scipy.fft import fft2, ifft2, fftshift, ifftshift
import matplotlib.pyplot as plt
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
import tkinter as tk
from tkinter import filedialog ,ttk
from concurrent.futures import ThreadPoolExecutor


# 设置中文字体显示
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

class ImageProcessingApp:
    def __init__(self, root, image_path=None):
        # 初始化主窗口和图像处理参数
        self.root = root
        self.root.title('彩色图像处理')
        self.image = None          # 存储当前处理图像
        self.original_img = None   # 存储原始图像（修改初始化）
        self.filter_radius = 50    # 低通滤波器初始半径
        self.threads = 4           # 多线程处理的线程池大小

        # 创建GUI组件
        self.create_widgets()
        
        # 自动加载默认图像
        if image_path:
            self.load_image(image_path)  # 使用动态路径加载默认图像

    def create_widgets(self):
        # 创建界面控件（按钮、滑块、图像显示区）
        # 文件选择按钮
        self.btn_load = tk.Button(self.root, text='选择图像', command=self.load_image)
        self.btn_load.grid(row=0, column=0, padx=10, pady=5)

        # 滤波器半径调节滑块
        self.lbl_radius = tk.Label(self.root, text='低通滤波器半径:')
        self.lbl_radius.grid(row=0, column=1, padx=10, pady=5)
        self.scale_radius = tk.Scale(self.root, from_=10, to=200, 
                                   orient=tk.HORIZONTAL, command=self.update_radius)
        self.scale_radius.set(50)
        self.scale_radius.grid(row=0, column=2, padx=10, pady=5)

        # 使用matplotlib创建图像显示画布
        self.fig, self.axes = plt.subplots(2, 3, figsize=(12, 8))
        self.canvas = FigureCanvasTkAgg(self.fig, master=self.root)

        # 修改画布布局行号
        self.canvas.get_tk_widget().grid(row=2, column=0, columnspan=3, padx=10, pady=10)

    def load_image(self, file_path=None):
        # 文件对话框选择图像文件
        if file_path is None:
            file_path = filedialog.askopenfilename(filetypes=[('图像文件', '*.bmp;*.png;*.jpg')])
        if file_path:
            # 使用imdecode处理包含中文路径的文件
            try:
                with open(file_path, 'rb') as f:
                    img_data = np.frombuffer(f.read(), dtype=np.uint8)
                    # 以彩色模式读取图像
                    self.original_img = cv2.imdecode(img_data, cv2.IMREAD_COLOR)
            except Exception as e:
                print("图像读取失败:", e)
                self.original_img = None
            
            # 加载后立即处理显示
            if self.original_img is not None:
                # 创建并初始化进度条
                self.progress = ttk.Progressbar(
                    self.root, 
                    length=400, 
                    mode='determinate',
                    maximum=100,
                    value=0
                )
                self.progress.grid(row=1, column=0, columnspan=3, padx=10, pady=5)
                self.process_and_display()

    def process_and_display(self):
        # 阶段1：图像预处理
        for ax in self.axes.flat:
            ax.cla()
        self.progress['value'] = 20
        self.progress.update()

        # 显示原始彩色图像
        self.axes[0, 0].imshow(cv2.cvtColor(self.original_img, cv2.COLOR_BGR2RGB))
        self.axes[0, 0].set_title('原始图像')
        self.axes[0, 0].axis('off')
        self.progress['value'] = 30
        self.progress.update()

        # 阶段2：通道处理
        with ThreadPoolExecutor(max_workers=self.threads) as executor:
            channels = cv2.split(self.original_img)
            futures = []
            for channel in channels:
                futures.append(executor.submit(self.process_channel, channel))
            
            results = []
            for future in futures:
                results.append(future.result())
            
            fshifts, low_pass_results, high_pass_results = zip(*results)

            # 合并处理结果
            fshift = cv2.merge(fshifts)
            img_low_pass = cv2.merge(low_pass_results)
            img_high_pass = cv2.merge(high_pass_results)
            self.progress['value'] = 50
            self.progress.update()

        # 阶段3：结果显示
        magnitude_spectrum = 20 * np.log(np.abs(fshift).clip(1e-6))
        self.axes[0, 1].imshow(magnitude_spectrum.astype(np.uint8), cmap='gray')
        self.axes[0, 1].set_title('幅度谱')
        self.axes[0, 1].axis('off')

        img_low_pass = img_low_pass.astype(np.uint8)
        self.axes[0, 2].imshow(cv2.cvtColor(img_low_pass, cv2.COLOR_BGR2RGB))
        self.axes[0, 2].set_title(f'低通滤波（半径{self.filter_radius}）')
        self.axes[0, 2].axis('off')

        img_high_pass = img_high_pass.astype(np.uint8)
        self.axes[1, 0].imshow(cv2.cvtColor(img_high_pass, cv2.COLOR_BGR2RGB))
        self.axes[1, 0].set_title(f'高通滤波（半径{self.filter_radius}）')
        self.axes[1, 0].axis('off')

        # Sobel处理
        gray_img = cv2.cvtColor(self.original_img, cv2.COLOR_BGR2GRAY)
        sobelx = cv2.Sobel(gray_img, cv2.CV_64F, 1, 0, ksize=5)
        sobely = cv2.Sobel(gray_img, cv2.CV_64F, 0, 1, ksize=5)
        self.progress['value'] = 80
        self.progress.update()

        self.axes[1, 1].imshow(sobelx, cmap='gray')
        self.axes[1, 1].set_title('Sobel X方向')
        self.axes[1, 1].axis('off')
        self.axes[1, 2].imshow(sobely, cmap='gray')
        self.axes[1, 2].set_title('Sobel Y方向')
        self.axes[1, 2].axis('off')

        # 完成处理
        self.progress['value'] = 100
        self.progress.update()
        
        # 隐藏进度条
        self.progress.grid_remove()
        self.canvas.draw()

    def update_radius(self, value):
        # 更新低通滤波器半径参数
        self.filter_radius = int(value)
        if self.original_img is not None:
            # 参数变化后重新处理图像
            self.process_and_display()

    def process_channel(self, channel):
        # 单通道图像处理流程
        # 二维傅里叶变换
        f = fft2(channel)
        # 将零频率分量移到频谱中心
        fshift = fftshift(f)
        
        # 创建低通滤波器掩膜
        rows, cols = channel.shape
        crow, ccol = rows // 2, cols // 2
        low_pass_filter = np.zeros((rows, cols), np.uint8)
        # 绘制圆形滤波器区域
        cv2.circle(low_pass_filter, (ccol, crow), self.filter_radius, 1, -1)
        
        # 应用低通滤波器
        f_low_pass = fshift * low_pass_filter
        # 逆傅里叶变换还原图像
        img_low_pass = np.abs(ifft2(ifftshift(f_low_pass)))
        
        # 创建高通滤波器（通过取反低通滤波器）
        high_pass_filter = np.ones((rows, cols), np.uint8) - low_pass_filter
        # 应用高通滤波器
        f_high_pass = fshift * high_pass_filter
        # 逆变换获得高通结果
        img_high_pass = np.abs(ifft2(ifftshift(f_high_pass)))
        
        # 返回傅里叶变换的幅值谱、低通和高通结果
        # 注意：返回fshift的幅值而非复数本身,无法加载出虚数部分
        return np.abs(fshift), img_low_pass, img_high_pass

if __name__ == '__main__':
    # 动态获取资源路径（用于打包后访问house.bmp）
    if getattr(sys, 'frozen', False):
        # 如果是打包后的环境，使用_MEIPASS路径
        base_path = sys._MEIPASS
    else:
        base_path = os.path.abspath('.')
    
    default_image_path = os.path.join(base_path, 'house.bmp')

    root = tk.Tk()
    app = ImageProcessingApp(root, default_image_path)  # 传递默认图片路径参数
    root.mainloop()